﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Raid6
{
    class Program
    {
        static void Main(string[] args)
        {
            var raid6Array = new Raid6Array(10, 4);
            Console.WriteLine("Original Data");
            Console.WriteLine(raid6Array);

            Console.WriteLine("Fail single data drive");
            raid6Array.DataDisks[0].FailDrive();
            Console.WriteLine(raid6Array);
            Console.WriteLine("Recover data");
            raid6Array.RebuildArray();
            Console.WriteLine(raid6Array);

            Console.WriteLine("Fail single data drive and Q drive");
            raid6Array.DataDisks[0].FailDrive();
            raid6Array.QDisk.FailDrive();
            Console.WriteLine(raid6Array);
            Console.WriteLine("Recover data");
            raid6Array.RebuildArray();
            Console.WriteLine(raid6Array);

            Console.WriteLine("Fail P drive and Q drive");
            raid6Array.PDisk.FailDrive();
            raid6Array.QDisk.FailDrive();
            Console.WriteLine(raid6Array);
            Console.WriteLine("Recover data");
            raid6Array.RebuildArray();
            Console.WriteLine(raid6Array);

            Console.WriteLine("Fail single Data and P drive");
            raid6Array.DataDisks[0].FailDrive();
            raid6Array.PDisk.FailDrive();
            Console.WriteLine(raid6Array);
            Console.WriteLine("Recover data");
            raid6Array.RebuildArray();
            Console.WriteLine(raid6Array);

            Console.WriteLine("Fail 2 Data");
            raid6Array.DataDisks[0].FailDrive();
            raid6Array.DataDisks[1].FailDrive();
            Console.WriteLine(raid6Array);
            Console.WriteLine("Recover data");
            raid6Array.RebuildArray();
            Console.WriteLine(raid6Array);

            Console.ReadLine();
        }

        public class Raid6Array
        {
            private readonly int _dataDiskCount;
            private readonly int _diskSize;
            private static readonly GaliosFieldElement GaliosFieldGenerator = new GaliosFieldElement(0x02);

            public readonly List<DiskDrive> DataDisks = new List<DiskDrive>();
            public readonly DiskDrive PDisk;
            public readonly DiskDrive QDisk;

            public Raid6Array(int diskSize, int dataDiskCount)
            {
                _dataDiskCount = dataDiskCount;
                _diskSize = diskSize;

                //Initialize data disks with random data
                for (var i = 0; i < _dataDiskCount; i++)
                {
                    DataDisks.Add(new DiskDrive(_diskSize, i));
                }

                //Initialize parity disks
                PDisk = new DiskDrive(_diskSize, 0);
                QDisk = new DiskDrive(_diskSize, 0);

                //Calculate P Values
                for (var i = 0; i < _diskSize; i++)
                {
                    PDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, DataDisks));
                }

                //Calculate Q Values
                for (var i = 0; i < _diskSize; i++)
                {
                    QDisk.DataBytes[i] = CalculateQValue(GetDriveBytesAtIndex(i, DataDisks));
                }
            }

            private static byte[] GetDriveBytesAtIndex(int index, List<DiskDrive> diskDrives)
            {
                var bytesAtIndex = new byte[diskDrives.Count];
                for (var i = 0; i < diskDrives.Count; i++)
                {
                    bytesAtIndex[i] = diskDrives[i].DataBytes[index];
                }
                return bytesAtIndex;
            }

            private static byte CalculatePValue(IList<byte> dataBytes)
            {
                if (dataBytes.Count() < 2)
                {
                    throw new Exception("Not enough drives to compute");
                }

                var pValue = dataBytes[0] ^ dataBytes[1];

                for (var i = 2; i < dataBytes.Count(); i++)
                {
                    pValue = pValue ^ dataBytes[i];
                }

                return (byte)pValue;
            }

            private static byte CalculateQValue(IList<byte> dataBytes)
            {
                if (dataBytes.Count() < 2)
                {
                    throw new Exception("Not enough drives to compute");
                }

                var qValue = new GaliosFieldElement(dataBytes[0]) + (new GaliosFieldElement(dataBytes[1]) * GaliosFieldGenerator);

                for (var i = 2; i < dataBytes.Count(); i++)
                {
                    var temp = new GaliosFieldElement(dataBytes[i]);
                    for (var j = 0; j < i; j++)
                    {
                        temp *= GaliosFieldGenerator;
                    }
                }

                return qValue.Value;
            }

            public class DiskDrive
            {
                public int DiskIndex { get; private set; }

                public byte[] DataBytes { get; private set; }

                private static readonly Random Random = new Random();

                public DiskDrive(int size, int diskIndex)
                {
                    DiskIndex = diskIndex;
                    DataBytes = new byte[size];
                    for (var i = 0; i < size; i++)
                    {
                        DataBytes[i] = (byte)Random.Next();
                    }
                }

                public override string ToString()
                {
                    return BitConverter.ToString(DataBytes);
                }

                public bool DriveFailed { get; private set; }

                public void FailDrive()
                {
                    DriveFailed = true;
                    for (var i = 0; i < DataBytes.Count(); i++)
                    {
                        DataBytes[i] = 0x00;
                    }
                }

                public void SetDriveHealthy()
                {
                    DriveFailed = false;
                }
            }

            public override string ToString()
            {
                var sb = new StringBuilder();
                sb.AppendLine("Data State:");

                for (int i = 0; i < _dataDiskCount; i++)
                {
                    sb.AppendLine("Data Disk " + i + ":" + DataDisks[i]);
                }

                sb.AppendLine("P Disk:" + PDisk);
                sb.AppendLine("Q Disk:" + QDisk);
                return sb.ToString();
            }

            public void RebuildArray()
            {
                var qDriveFailed = QDisk.DriveFailed;
                var pDriveFailed = PDisk.DriveFailed;
                var failedDataDiskCount = (from m in DataDisks
                                           where m.DriveFailed
                                           select m).Count();
                if (qDriveFailed && !pDriveFailed && failedDataDiskCount == 0)
                {
                    //Recover Q drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        QDisk.DataBytes[i] = CalculateQValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    QDisk.SetDriveHealthy();
                    return;
                }
                if (!qDriveFailed && pDriveFailed && failedDataDiskCount == 0)
                {
                    //Recover P drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        PDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    PDisk.SetDriveHealthy();
                    return;
                }
                if (!qDriveFailed && !pDriveFailed && failedDataDiskCount == 1)
                {
                    //Recover failed data disk
                    var recoveryDataDisks = new List<DiskDrive> { PDisk };
                    recoveryDataDisks.AddRange(from m in DataDisks
                                               where !m.DriveFailed
                                               select m);
                    var failedDataDisk = (from m in DataDisks
                                          where m.DriveFailed
                                          select m).First();
                    for (var i = 0; i < _diskSize; i++)
                    {
                        failedDataDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, recoveryDataDisks));
                    }
                    failedDataDisk.SetDriveHealthy();
                    return;
                }
                if (qDriveFailed && !pDriveFailed && failedDataDiskCount == 1)
                {
                    //Recover failed data disk
                    var recoveryDataDisks = new List<DiskDrive> { PDisk };
                    recoveryDataDisks.AddRange(from m in DataDisks
                                               where !m.DriveFailed
                                               select m);
                    var failedDataDisk = (from m in DataDisks
                                          where m.DriveFailed
                                          select m).First();
                    for (var i = 0; i < _diskSize; i++)
                    {
                        failedDataDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, recoveryDataDisks));
                    }
                    failedDataDisk.SetDriveHealthy();

                    //Recover Q drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        QDisk.DataBytes[i] = CalculateQValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    QDisk.SetDriveHealthy();
                    return;
                }
                if (!qDriveFailed && pDriveFailed && failedDataDiskCount == 1)
                {
                    //Recover data drive
                    var diskIndex = (from m in DataDisks
                                     where m.DriveFailed
                                     select m.DiskIndex).First();
                    var failedDataDisk = (from m in DataDisks
                                          where m.DriveFailed
                                          select m).First();
                    for (var i = 0; i < _diskSize; i++)
                    {
                        var qx = new GaliosFieldElement(CalculateQValue(GetDriveBytesAtIndex(i, DataDisks)));

                        failedDataDisk.DataBytes[i] =
                            (GaliosFieldMath.Pow(GaliosFieldGenerator, 255 - diskIndex)*
                             (qx + new GaliosFieldElement(QDisk.DataBytes[i]))).Value;
                    }
                    failedDataDisk.SetDriveHealthy();

                    //Recover P drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        PDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    PDisk.SetDriveHealthy();
                    return;
                }
                if (qDriveFailed && pDriveFailed && failedDataDiskCount == 0)
                {
                    //Recover P drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        PDisk.DataBytes[i] = CalculatePValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    PDisk.SetDriveHealthy();
                    //Recover Q drive
                    for (var i = 0; i < _diskSize; i++)
                    {
                        QDisk.DataBytes[i] = CalculateQValue(GetDriveBytesAtIndex(i, DataDisks));
                    }
                    QDisk.SetDriveHealthy();
                    return;
                }
                if (!qDriveFailed && !pDriveFailed && failedDataDiskCount == 2)
                {
                    //Recover data drive
                    var diskIndexX = (from m in DataDisks
                                      where m.DriveFailed
                                      select m.DiskIndex).First();
                    var failedDataDiskX = (from m in DataDisks
                                           where m.DriveFailed
                                           select m).First();
                    var diskIndexY = (from m in DataDisks
                                      where m.DriveFailed
                                      select m.DiskIndex).Last();
                    var failedDataDiskY = (from m in DataDisks
                                           where m.DriveFailed
                                           select m).Last();
                    for (var i = 0; i < _diskSize; i++)
                    {
                        var qxy = new GaliosFieldElement(CalculateQValue(GetDriveBytesAtIndex(i, DataDisks)));
                        var pxy = new GaliosFieldElement(CalculatePValue(GetDriveBytesAtIndex(i, DataDisks)));
                        var pDiskGalios = new GaliosFieldElement(PDisk.DataBytes[i]);
                        var qDiskGalios = new GaliosFieldElement(QDisk.DataBytes[i]);

                        var galiosYminusX = GaliosFieldMath.Pow(GaliosFieldGenerator, diskIndexY - diskIndexX);
                        var galiosNegativeX = GaliosFieldMath.Pow(GaliosFieldGenerator, 255 - diskIndexX);
                        var galiosYminusXxorIdentInverse = GaliosFieldMath.Pow(galiosYminusX + new GaliosFieldElement(0x01), 254);
                        var tempA = galiosYminusX * galiosYminusXxorIdentInverse;
                        var tempB = galiosNegativeX * galiosYminusXxorIdentInverse;
                        var xDiskGalios = (tempA*(pDiskGalios + pxy)) + (tempB*(qDiskGalios + qxy));
                        var yDiskGalios = (pDiskGalios + pxy) + xDiskGalios;
                        failedDataDiskX.DataBytes[i] = xDiskGalios.Value;
                        failedDataDiskY.DataBytes[i] = yDiskGalios.Value;
                    }
                    failedDataDiskX.SetDriveHealthy();
                    failedDataDiskY.SetDriveHealthy();
                    return;
                }
                throw new Exception("Too many drives are failed to recover data");
            }
        }

        public class GaliosFieldElement
        {
            public byte Value { get; set; }

            public GaliosFieldElement(byte value)
            {
                Value = value;
            }

            public static GaliosFieldElement operator +(GaliosFieldElement a, GaliosFieldElement b)
            {
                return new GaliosFieldElement((byte)(a.Value ^ b.Value));
            }

            public static GaliosFieldElement operator *(GaliosFieldElement a, GaliosFieldElement b)
            {
                var aBitArray = new BitArray(new[] { a.Value });
                GaliosFieldElement result = null;

                for (var i = 0; i < 8; i++)
                {
                    var ai = aBitArray[i];
                    var b2I = b;
                    for (int j = 0; j < i; j++)
                    {
                        b2I = GaliosFieldMath.MultiplyByGalois02(b2I);
                    }

                    var aib2I = ai ? b2I : new GaliosFieldElement(0x00);
                    if (result != null)
                    {
                        result = result + aib2I;
                    }
                    else
                    {
                        result = aib2I;
                    }
                }

                return result;
            }
        }

        public static class GaliosFieldMath
        {
            public static GaliosFieldElement Pow(GaliosFieldElement gfe, int power)
            {
                while (power < 0)
                {
                    power += 255;
                }
                while (power > 255)
                {
                    power -= 255;
                }
                if (power == 0)
                {
                    return new GaliosFieldElement(0x01);
                }
                var result = gfe;
                for (var i = 1; i < power; i++)
                {
                    result *= gfe;
                }
                return result;
            }

            public static GaliosFieldElement MultiplyByGalois02(GaliosFieldElement gfe)
            {
                return new GaliosFieldElement((byte)((gfe.Value << 1) ^ ((gfe.Value & (byte)0x80) != 0 ? 0x1d : 0)));
            }
        }
    }
}
